package intcomp

import (
	"encoding/binary"

	"gitee.com/haodreams/libs/easy"
)

func ZigZagEncode(value int64) uint64 {
	return uint64((value << 1) ^ (value >> 63))
}

func ZigZagDecode(encoded uint64) int64 {
	return int64((encoded >> 1) ^ (-(encoded & 1)))
}

func ZigZagEncode32(value int32) uint32 {
	return uint32((value << 1) ^ (value >> 31))
}

func ZigZagDecode32(encoded uint32) int32 {
	return int32((encoded >> 1) ^ (-(encoded & 1)))
}

// 整型压缩算法，移除0的字节
// 最高3位表示使用的字节数
// 最高位为1 表示 最高的三位是后面数据包的字节数
// 100 后续一个字节
// 111 后续7个字节
type IntComp struct {
	cap  int
	size int
	buf  *easy.Buffer
	bs   []byte
}

func NewIntComp(size int) *IntComp {
	m := new(IntComp)
	m.cap = size
	m.buf = easy.NewBufferWithSize(m.cap * 8)
	m.bs = make([]byte, 8)
	m.buf.Write(m.bs[:4]) //空出前4个字节
	return m
}

func (m *IntComp) Write(ival int64) (err error) {
	uval := ZigZagEncode(ival)
	switch {
	case uval < 0x1f:
		err = m.buf.WriteByte(byte(uval))
	case uval < 0x1fff:
		m.bs[0] = byte(uval >> 8)
		m.bs[1] = byte(uval)
		m.bs[0] |= 1 << 5
		_, err = m.buf.Write(m.bs[:2])
	case uval < 0x1fffff:
		m.bs[0] = byte(uval >> 16)
		m.bs[1] = byte(uval >> 8)
		m.bs[2] = byte(uval)
		m.bs[0] |= 2 << 5
		_, err = m.buf.Write(m.bs[:3])
	case uval < 0x1fffffff:
		m.bs[0] = byte(uval >> 24)
		m.bs[1] = byte(uval >> 16)
		m.bs[2] = byte(uval >> 8)
		m.bs[3] = byte(uval)
		m.bs[0] |= 3 << 5
		_, err = m.buf.Write(m.bs[:4])
	case uval < 0x1fffffffff:
		m.bs[0] = byte(uval >> 32)
		m.bs[1] = byte(uval >> 24)
		m.bs[2] = byte(uval >> 16)
		m.bs[3] = byte(uval >> 8)
		m.bs[4] = byte(uval)
		m.bs[0] |= 4 << 5
		_, err = m.buf.Write(m.bs[:5])
	case uval < 0x1fffffffffff:
		m.bs[0] = byte(uval >> 40)
		m.bs[1] = byte(uval >> 32)
		m.bs[2] = byte(uval >> 24)
		m.bs[3] = byte(uval >> 16)
		m.bs[4] = byte(uval >> 8)
		m.bs[5] = byte(uval)
		m.bs[0] |= 5 << 5
		_, err = m.buf.Write(m.bs[:6])
	case uval < 0x1fffffffffffff:
		m.bs[0] = byte(uval >> 48)
		m.bs[1] = byte(uval >> 40)
		m.bs[2] = byte(uval >> 32)
		m.bs[3] = byte(uval >> 24)
		m.bs[4] = byte(uval >> 16)
		m.bs[5] = byte(uval >> 8)
		m.bs[6] = byte(uval)
		m.bs[0] |= 6 << 5
		_, err = m.buf.Write(m.bs[:7])
	case uval < 0x1fffffffffffffff:
		m.bs[0] = byte(uval >> 56)
		m.bs[1] = byte(uval >> 48)
		m.bs[2] = byte(uval >> 40)
		m.bs[3] = byte(uval >> 32)
		m.bs[4] = byte(uval >> 24)
		m.bs[5] = byte(uval >> 16)
		m.bs[6] = byte(uval >> 8)
		m.bs[7] = byte(uval)
		m.bs[0] |= 7 << 5
		_, err = m.buf.Write(m.bs[:8])
	default: //数据无效
		err = m.buf.WriteByte(0)
	}
	m.size++
	return
}

func (m *IntComp) Flush() []byte {
	binary.BigEndian.PutUint32(m.bs, uint32(m.size))
	m.buf.WriteAt(m.bs[:4], 0)
	return m.buf.Bytes()
}

func (m *IntComp) Decode(bs []byte) (ivals []int64, err error) {
	data := make([]byte, 8)
	buf := easy.NewBuffer(bs)
	buf.SetBigEndian(true)
	size := int(buf.Int32())
	ivals = make([]int64, size)
	for i := 0; i < size; i++ {
		c := buf.Byte()
		num := c >> 5
		val := c & 0x1f
		switch num {
		case 0:
			ivals[i] = ZigZagDecode(uint64(val))
		case 1:
			v := buf.Byte()
			uval := uint64(val)<<8 | uint64(v)
			ivals[i] = ZigZagDecode(uval)
		case 2:
			buf.Read(data[:2])
			uval := uint64(val)<<16 | uint64(data[0])<<8 | uint64(data[1])
			ivals[i] = ZigZagDecode(uval)
		case 3:
			buf.Read(data[:3])
			uval := uint64(val)<<24 | uint64(data[0])<<16 | uint64(data[1])<<8 | uint64(data[2])
			ivals[i] = ZigZagDecode(uval)
		case 4:
			buf.Read(data[:4])
			uval := uint64(val)<<32 | uint64(data[0])<<24 | uint64(data[1])<<16 | uint64(data[2])<<8 | uint64(data[3])
			ivals[i] = ZigZagDecode(uval)
		case 5:
			buf.Read(data[:5])
			uval := uint64(val)<<40 | uint64(data[0])<<32 | uint64(data[1])<<24 | uint64(data[2])<<16 | uint64(data[3])<<8 | uint64(data[4])
			ivals[i] = ZigZagDecode(uval)
		case 6:
			buf.Read(data[:6])
			uval := uint64(val)<<48 | uint64(data[0])<<40 | uint64(data[1])<<32 | uint64(data[2])<<24 | uint64(data[3])<<16 | uint64(data[4])<<8 | uint64(data[5])
			ivals[i] = ZigZagDecode(uval)
		case 7:
			buf.Read(data[:7])
			uval := uint64(val)<<56 | uint64(data[0])<<48 | uint64(data[1])<<40 | uint64(data[2])<<32 | uint64(data[3])<<24 | uint64(data[4])<<16 | uint64(data[5])<<8 | uint64(data[6])
			ivals[i] = ZigZagDecode(uval)
		}
	}
	return
}

func (m *IntComp) Resize(size int) {
	if size <= m.cap {
		m.Reset()
		return
	}
	m.cap = size
	m.buf = easy.NewBufferWithSize(m.cap * 8)
	m.Reset()
}

func (m *IntComp) Reset() {
	m.size = 0
	m.buf.Reset()
	m.buf.Write(m.bs[:4]) //空出前4个字节
}
